#include <iostream>
#include <Eigen/Dense>
#include <vector>
#include <unordered_map>

using namespace Eigen;
using namespace std;

const int ROBOT_DIM = 3;
const int LANDMARK_DIM = 2;

struct Landmark {
    Vector2d position;      
    Vector2d first_estimate; // FEJ
    bool is_initialized;    
};

class EKFSLAM {
public:
    EKFSLAM() {
        robot_state_ = Vector3d(0, 0, 0);
        covariance_ = MatrixXd::Identity(ROBOT_DIM, ROBOT_DIM) * 0.1;
    }

    void ProcessMeasurement(int landmark_id, double r, double phi) {
        if (landmarks_.find(landmark_id) == landmarks_.end() || !landmarks_[landmark_id].is_initialized) {
            // 首次观测到路标点，初始化并保存FEJ点
            InitializeLandmark(landmark_id, r, phi);
        } else {
            // 已初始化的路标点，使用FEJ点计算雅可比
            UpdateLandmark(landmark_id, r, phi);
        }
    }

    // 打印状态
    void PrintState() {
        cout << "Robot State: [" << robot_state_.transpose() << "]" << endl;
        for (auto& lm : landmarks_) {
            cout << "Landmark " << lm.first << ": [" << lm.second.position.transpose() << "] (FEJ: [" 
                 << lm.second.first_estimate.transpose() << "])" << endl;
        }
    }

    Vector3d robot_state_;                          // 机器人状态
    MatrixXd covariance_;                          // 协方差矩阵
    unordered_map<int, Landmark> landmarks_;       // 路标点集合

private:

    void InitializeLandmark(int id, double r, double phi) {
        double x = robot_state_(0) + r * cos(robot_state_(2) + phi);
        double y = robot_state_(1) + r * sin(robot_state_(2) + phi);

        // cout << "x ======= : " << x << endl;
        // cout << "y ======= : " << y << endl;

        landmarks_[id] = {Vector2d(x, y), Vector2d(x, y), true};

        // extern covariance
        int old_size = covariance_.rows();
        covariance_.conservativeResize(old_size + LANDMARK_DIM, old_size + LANDMARK_DIM);
        covariance_.block(old_size, 0, LANDMARK_DIM, old_size).setZero();
        covariance_.block(0, old_size, old_size, LANDMARK_DIM).setZero();
        covariance_.block(old_size, old_size, LANDMARK_DIM, LANDMARK_DIM) = MatrixXd::Identity(LANDMARK_DIM, LANDMARK_DIM) * 1e6;
        
        // cout << "covariance_ ======= : \n" << covariance_ << endl;

    }

    void UpdateLandmark(int id, double r, double phi) {
        Landmark& lm = landmarks_[id];

        // observer model h = [distance; bearing]
        Vector2d z_pred;
        // 不使用FEJ的话，用lm.position
        Vector2d delta = lm.first_estimate - robot_state_.head(2); // FEJ点处计算, 
        z_pred(0) = delta.norm();                               // 预测距离
        z_pred(1) = atan2(delta(1), delta(0)) - robot_state_(2);  // 预测方位角

        // 观测残差
        Vector2d z_meas(r, phi);
        Vector2d residual = z_meas - z_pred;

        // 计算雅可比矩阵 H（在FEJ点处）
        MatrixXd H = MatrixXd::Zero(2, ROBOT_DIM + landmarks_.size() * LANDMARK_DIM);
        // delta （在FEJ点处）
        double q = delta.squaredNorm();
        double sqrt_q = sqrt(q);
        H.block(0, 0, 1, 2) = -delta.transpose() / sqrt_q;                // 对机器人x,y的导数
        H(0, 2) = 0;                                                     // 对机器人theta的导数
        H.block(1, 0, 1, 2) = Vector2d(delta(1), -delta(0)).transpose() / q; // 对机器人x,y的导数
        H(1, 2) = -1;                                                     // 对机器人theta的导数

        // 路标点的雅可比（在FEJ点处）
        int lm_index = ROBOT_DIM + distance(landmarks_.begin(), landmarks_.find(id)) * LANDMARK_DIM;
        H.block(0, lm_index, 1, 2) = delta.transpose() / sqrt_q;           // 对路标点的导数
        H.block(1, lm_index, 1, 2) = Vector2d(-delta(1), delta(0)).transpose() / q;

        // 观测噪声矩阵
        Matrix2d R = Matrix2d::Identity() * 0.01;

        // EKF更新步骤
        MatrixXd K = covariance_ * H.transpose() * (H * covariance_ * H.transpose() + R).inverse();
        VectorXd state_update = K * residual;
        robot_state_ += state_update.head(ROBOT_DIM);
        for (auto& lm_pair : landmarks_) {
            int idx = ROBOT_DIM + distance(landmarks_.begin(), landmarks_.find(lm_pair.first)) * LANDMARK_DIM;
            lm_pair.second.position += state_update.segment(idx, LANDMARK_DIM);
        }
        covariance_ = (MatrixXd::Identity(covariance_.rows(), covariance_.cols()) - K * H) * covariance_;

        // cout << "covariance_ ======= : \n" << covariance_ << endl;
    }
};

int main() {
    EKFSLAM ekf;

    // 模拟数据：机器人移动并观测路标点
    ekf.ProcessMeasurement(1, 1.0, 0.0);      // 观测到路标点1（距离1，方位角0）
    ekf.robot_state_ = Vector3d(0.5, 0.0, 0.0); // 机器人移动
    ekf.ProcessMeasurement(1, 0.5, 0.0);      // 再次观测路标点1（使用FEJ）

    ekf.PrintState();
    return 0;
}